import numpy as np
from PIL import Image
import os.path as osp
import matplotlib.pyplot as plt
from tqdm import tqdm
def read_image(img_path):
    """Keep reading image until succeed.
    This can avoid IOError incurred by heavy IO process."""
    got_img = False
    if not osp.exists(img_path):
        raise IOError("{} does not exist".format(img_path))
    while not got_img:
        try:
            img = Image.open(img_path).convert('RGB')
            got_img = True
        except IOError:
            print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
            pass
    return img

def calculate_mean_std_in_dataset(dataset):
    """
    计算数据集RGB三通道的均值
    :param dataset:
    :return:
    """
    print("begin to calculate mean and std...........")
    means, stdevs = [], []
    img_list = []
    print("逐张读入全部图片")
    for (img_path, pid, camid) in tqdm(dataset):
        img = read_image(img_path)
        img = np.array(img)
        img = img[:, :, :, np.newaxis]
        img_list.append(img)
    imgs = np.concatenate(img_list, axis=3)
    imgs = imgs.astype(np.float32) / 255
    use_numpy = True  # False逐张计算，速度慢但占用内存少
    if use_numpy:
        print("numpy计算训练集的means和std")
        for i in tqdm(range(3)):
            pixels = imgs[:, :, i, :].ravel()
            means.append(np.mean(pixels))
            stdevs.append(np.std(pixels))
        print('means:{}'.format(means), 'std:{}'.format(stdevs))
    else:
        print("逐张计算means和std")
        R_channel, G_channel, B_channel = [], [], []
        for i in tqdm(range(imgs.shape[3])):
            pixelsR = imgs[:, :, 0, i].ravel()
            pixelsG = imgs[:, :, 1, i].ravel()
            pixelsB = imgs[:, :, 2, i].ravel()
            R_channel.append(np.mean(pixelsR))
            G_channel.append(np.mean(pixelsG))
            B_channel.append(np.mean(pixelsB))
        means = [np.mean(R_channel),np.mean(G_channel),np.mean(B_channel)]
        R_channel, G_channel, B_channel = [], [], []
        for i in tqdm(range(imgs.shape[3])):
            pixelsR = imgs[:, :, 0, i].ravel()
            pixelsG = imgs[:, :, 1, i].ravel()
            pixelsB = imgs[:, :, 2, i].ravel()
            R_channel.append(np.sum((pixelsR-means[0]) ** 2))
            G_channel.append(np.sum((pixelsG-means[1]) ** 2))
            B_channel.append(np.sum((pixelsB-means[2]) ** 2))
        image_size = imgs.shape[0]*imgs.shape[1]
        stdevs = [np.sqrt(np.mean(R_channel)/image_size), np.sqrt(np.mean(G_channel)/image_size),np.sqrt(np.mean(B_channel)/image_size)]
        print('means:{}'.format(means), 'std:{}'.format(stdevs))

    per_image_Rmean, per_image_Gmean, per_image_Bmean = [], [], []  # 记录每张图片的均值
    per_image_Rstd, per_image_Gstd, per_image_Bstd = [], [], []  # 记录每张图片的标准擦差
    for i in tqdm(range(imgs.shape[3])):
        pixelsR = imgs[:, :, 0, i].ravel()
        pixelsG = imgs[:, :, 1, i].ravel()
        pixelsB = imgs[:, :, 2, i].ravel()
        per_image_Rmean.append(np.mean(pixelsR))
        per_image_Gmean.append(np.mean(pixelsG))
        per_image_Bmean.append(np.mean(pixelsB))
        per_image_Rstd.append(np.std(pixelsR))
        per_image_Gstd.append(np.std(pixelsG))
        per_image_Bstd.append(np.std(pixelsB))
    print("统计RGB三通道均值和标准差的直方图")
    Rmean_statis_list, Gmean_statis_list, Bmean_statis_list = [], [], []
    Rstd_statis_list, Gstd_statis_list, Bstd_statis_list = [], [], []
    stride = 0.05  # 分割的步长
    for i in tqdm(np.arange(0, 1, stride)):
        per_image_Rmean = np.array(per_image_Rmean)
        Rmean_statis_list.append(sum((per_image_Rmean > i) & (per_image_Rmean <= i + stride)))
        per_image_Gmean = np.array(per_image_Gmean)
        Gmean_statis_list.append(sum((per_image_Gmean > i) & (per_image_Gmean <= i + stride)))
        per_image_Bmean = np.array(per_image_Bmean)
        Bmean_statis_list.append(sum((per_image_Bmean > i) & (per_image_Bmean <= i + stride)))

        per_image_Rstd = np.array(per_image_Rstd)
        Rstd_statis_list.append(sum((per_image_Rstd > i) & (per_image_Rstd <= i + stride)))
        per_image_Gstd = np.array(per_image_Gstd)
        Gstd_statis_list.append(sum((per_image_Gstd > i) & (per_image_Gstd <= i + stride)))
        per_image_Bstd = np.array(per_image_Bstd)
        Bstd_statis_list.append(sum((per_image_Bstd > i) & (per_image_Bstd <= i + stride)))
    plt.figure(1)
    plt.subplot(121)
    plt.plot(np.arange(0, 1, stride), Rmean_statis_list, label='R channel')
    plt.plot(np.arange(0, 1, stride), Gmean_statis_list, label='G channel')
    plt.plot(np.arange(0, 1, stride), Bmean_statis_list, label='B channel')
    plt.title('means=[{:.3f},{:.3f},{:.3f}]'.format(means[0],means[1],means[2]))
    plt.legend()  # 显示图例
    plt.subplot(122)
    plt.plot(np.arange(0, 1, stride), Rstd_statis_list, label='R channel')
    plt.plot(np.arange(0, 1, stride), Gstd_statis_list, label='G channel')
    plt.plot(np.arange(0, 1, stride), Bstd_statis_list, label='B channel')
    plt.title('stdevs=[{:.3f},{:.3f},{:.3f}]'.format(stdevs[0],stdevs[1],stdevs[2]))
    plt.legend()  # 显示图例
    plt.show()

    return means, stdevs

from collections import defaultdict
def calculate_pids_num(dataset):
    pids = defaultdict(list)
    for (img_path, pid, camid) in tqdm(dataset):
        pids[pid].append(img_path)
    num = []
    for i in pids.keys():
        num.append(len(pids[i]))
    num.sort()
    plt.plot(num)
    plt.show()
